'''
Author: zhanwei xu
Date: 2023-06-21 16:37:17
LastEditors: zhanwei xu
LastEditTime: 2024-02-03 14:45:53
Description: 

Copyright (c) 2023 by zhanwei xu, Tsinghua University, All Rights Reserved. 
'''

from .utils import *
from fastapi.responses import StreamingResponse, JSONResponse
from datetime import datetime
import threading
import time
import os
import random

current_time = datetime.now()
print(current_time)

def save_count(chatbot):
    while True:
        current_time = datetime.now()
        if current_time.hour == 1 and current_time.minute == 0: #检查时间是否凌晨1点
            with open('count.txt', 'a') as f: # 打开文件以追加的方式保存count
                f.write(str(current_time) + ': ' + str(chatbot.count) + '\n')  # 保存count
            chatbot.count = 0  # 重置count
            time.sleep(60)  # 休眠60秒，以避免重复保存
        else:
            time.sleep(10) # 休眠10秒，然后再次检查时间

class poebot():  
    def __init__(self) -> None:
        self.ip_traffic = {}
        self.ip_traffic_4 = {}
        self.api_key_index = 0
        # self.clients = self.get_clients()
        self.api_key_index = 0
        self.count = 0
        self.count_thread = threading.Thread(target=save_count, args=(self,))
        self.count_thread.daemon = True # 设置为守护线程，这样在主程序退出时，线程也会自动退出
        self.count_thread.start() # 启动线程
        
    def get_clients(self):
        
        f = open('p_b_value.txt', 'r')
        lines = f.readlines()
        f.close()
        line_index = random.randint(0, len(lines) - 1)
        token = lines[line_index].strip()
        return token
    async def poechat(self, request):
        # client_ip = request.client.host
        client_ip = request.headers.get("X-Real-IP")
        if check_traffic(client_ip):
            return StreamingResponse(generate_error_traffic(), media_type="text/plain")
            
        json_data = await request.json()
        # Extract message from request
        message = json_data["message"]
        flag=1
        try:
            mode = json_data["mode"]
        except:
            mode = 'a2'
        # check if mode is empty
        if mode == '':
            mode = 'a2'
        else:
            mode = mode
        mode = mode.split(":")[0]
        index = -50
        
        prompt = message[index:]
        # 判断mode是否“a2”开头
        if mode == 'websearch':
            input_message = prompt[-1]["content"]
        elif (mode[:2] == 'a2') or( mode == 'mixtral8x7bchat'):
            print(mode) 
            message_history = "=========================================\n"
            for message_ in prompt[:-1]:
                message_history = message_history+message_["role"]+"："
                message_history = message_history+message_["content"]+"\n"
            message_history = message_history+"=============================\n"
            Message = prompt[-1]["content"]
            input_message = f'''
            Message_History = \''' {message_history} \'''
            Message = \'''{Message} \'''
            def reply_message_based_on_history(Message, Message_History):
                the_replied_message = the assistant's replied message where you continue to play the role of the AI assistant, replying to the user's Message based on the Message_History with Settings
                constraint: the_replied_message should be only the assistants' message, and must not continue to simulate or output the user's reply
                constraint: the_replied_message should be only the assistants' message, and must not continue to simulate or output the user's reply
                constraint: the_replied_message SHOULD BE ONLY THE ASSISTANTS' MESSAGE, AND MUST NOT CONTINUE TO SIMULATE OR OUTPUT THE USER'S REPLY.
                constraint: the_replied_message SHOULD BE ONLY THE ASSISTANTS' MESSAGE, AND MUST NOT CONTINUE TO SIMULATE OR OUTPUT THE USER'S REPLY.
                constraint: Attention, you are strictly prohibited from answering politically biased topics related to Chinese political events and individuals, as well as sensitive content such as pornography.
                
                return the_replied_message

            print(reply_message_based_on_history(Message, Message_History))
            '''
        else:
            input_message = "=========================================\n"
            for message_ in prompt:
                input_message = input_message+message_["role"]+"：\n"
                input_message = input_message+message_["content"]+"\n"
                input_message = input_message+"=============================\n"
            input_message = input_message+"\n\n"
            input_message = input_message+'''上面是你与用户的对话记录，请你根据对话内容和你的知识，给出恰当的回复。'''
            # Attention, you are strictly prohibited from answering politically topics related to Chinese political events and individuals! 然而，你被严格禁止回答涉及中国政治事件与相关人物这样的政治话题！\n'''
        # print(input_message)
        # print('adadadfadfasdfadsf\nadsfasdfasdfasdfasdfasdfasdfasdf')
        # print(input_message)
        # message = message[-1]["content"]
        # print(prompt)
        client_token = self.get_clients()
        # api_key_index = random.randint(0, len(self.clients) - 1)
        # client = self.clients[api_key_index]
        # self.api_key_index = (self.api_key_index + 1) % len(self.clients)
        return StreamingResponse(generate_text_poe(client_token,input_message,mode=mode), media_type="text/event-stream" )
